using System.Collections.Generic;
using System.Runtime.CompilerServices;
using HTCollections;
using HTCollections.Unsafe;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

namespace GPUAnimationLib.Runtime
{
    /// <summary>
    /// 对象KdTree
    /// </summary>
    public class EntityKdTree
    {
        /// <summary>
        /// 允许屏幕外一定范围内的角色
        /// </summary>
        const float Tolerance = 1f;

        /// <summary>
        /// 是否启用
        /// </summary>
        public bool Enable { get; set; }

        /// <summary>
        /// 使用多少帧进行更新
        /// </summary>
        public int UpdateFrameCount { get; set; }

        /// <summary>
        /// 是否为有效KdTree
        /// </summary>
        public bool IsValidKdTree { get; private set; }

        /// <summary>
        /// 添加角色Transform缓冲
        /// </summary>
        Dictionary<int, Transform> mTransformTempAddCache;

        /// <summary>
        /// 移除角色Transform缓冲
        /// </summary>
        HashSet<int> mTransformTempRemoveCache;

        /// <summary>
        /// 角色Transform集合
        /// key: 角色动画Id
        /// value: 角色Transform
        /// </summary>
        Dictionary<int, Transform> mTransformCache;

        /// <summary>
        /// 角色动画Id集合
        /// </summary>
        List<int> mIdList;

        /// <summary>
        /// 在屏幕范围内的角色动画Id集合
        /// </summary>
        HashSet<int> mInAreaSet;

        /// <summary>
        /// 缓存角色坐标集合
        /// float3(float2(pos), float(id))
        /// </summary>
        UnsafeList<float3> mPositions;

        /// <summary>
        /// 角色坐标集合缓存的包裹器（用于在job计算中传递）
        /// </summary>
        JobDataWrap<UnsafeList<float3>> mPositionsWrap;

        /// <summary>
        /// KdTree空间划分使用的数据结构
        /// </summary>
        UnsafeKdTree2D mKdTree;

        /// <summary>
        /// KdTree的包裹器（用于在job计算中传递）
        /// </summary>
        JobDataWrap<UnsafeKdTree2D> mKdTreeWrap;

        /// <summary>
        /// 屏幕内角色坐标计算结果
        /// </summary>
        UnsafeList<float3> mResultBuffer;

        /// <summary>
        /// 屏幕内角色坐标计算结果包裹器（用于在job计算中传递）
        /// </summary>
        JobDataWrap<UnsafeList<float3>> mResultWrap;

        /// <summary>
        /// 剩余多少帧更新完毕
        /// </summary>
        int mUpdateFrameLeft;

        /// <summary>
        /// 每帧更新的角色数量
        /// </summary>
        int mPerFrameUpdateCount;

        /// <summary>
        /// 更新的开始索引
        /// </summary>
        int mUpdateIdBegin;

        /// <summary>
        /// 更新的结束索引
        /// </summary>
        int mUpdateIdEnd;

        /// <summary>
        /// KdTree构建任务句柄
        /// </summary>
        JobHandle mBuildJobHandle;

        /// <summary>
        /// 屏幕内角色查询任务句柄
        /// </summary>
        JobHandle mQueryJobHandle;

        /// <summary>
        /// KdTree构建任务可以完成
        /// </summary>
        bool mBuildJobCanToComplete;

        /// <summary>
        /// 屏幕内角色查询任务可以完成
        /// </summary>
        bool mQueryJobCanToComplete;

        bool mIsValid;

        public EntityKdTree()
        {
            mIsValid = true;
            mTransformTempAddCache = new Dictionary<int, Transform>(200);
            mTransformTempRemoveCache = new HashSet<int>();
            mTransformCache = new Dictionary<int, Transform>(200);
            mIdList = new List<int>(200);
            mInAreaSet = new HashSet<int>();
            mPositions = new UnsafeList<float3>(200, Allocator.Persistent);
            mPositionsWrap = new JobDataWrap<UnsafeList<float3>>(mPositions, Allocator.Persistent);
            mKdTree = new UnsafeKdTree2D(8, Allocator.Persistent);
            mKdTreeWrap = new JobDataWrap<UnsafeKdTree2D>(mKdTree, Allocator.Persistent);
            mResultBuffer = new UnsafeList<float3>(200, Allocator.Persistent);
            mResultWrap = new JobDataWrap<UnsafeList<float3>>(mResultBuffer, Allocator.Persistent);
            Enable = true;
            UpdateFrameCount = 10;
            mUpdateFrameLeft = UpdateFrameCount;
            mUpdateIdBegin = -1;
            mBuildJobCanToComplete = false;
            mQueryJobCanToComplete = false;
            IsValidKdTree = false;
        }

        public void Dispose()
        {
            if (!mIsValid) return;
            mIsValid = false;
            mTransformTempAddCache = null;
            mTransformTempRemoveCache = null;
            mTransformCache = null;
            mIdList = null;
            mInAreaSet = null;
            if (mPositions.IsCreated)
                mPositions.Dispose();
            mPositionsWrap.Dispose();
            mKdTree.Dispose();
            mKdTreeWrap.Dispose();
            if (mResultBuffer.IsCreated)
                mResultBuffer.Dispose();
            mResultWrap.Dispose();
        }

        /// <summary>
        /// 同步KdTree构建
        /// 发起屏幕角色查询任务
        /// </summary>
        public void Update1()
        {
            if (!mIsValid || !Enable || UpdateFrameCount <= 0) return;

            if (mBuildJobCanToComplete)
            {
                //完成KdTree构建任务
                mBuildJobCanToComplete = false;
                mBuildJobHandle.Complete();
                mKdTree = mKdTreeWrap.Value;
                IsValidKdTree = true;
                if (CalculateArea(out var rect))
                {
                    //开始查找屏幕内角色
                    mResultBuffer.Clear();
                    mResultWrap.Value = mResultBuffer;
                    KdTreeQueryRectJob queryJob = new KdTreeQueryRectJob()
                    {
                        KdTree = mKdTree,
                        Rect = rect,
                        Result = mResultWrap,
                    };
                    mQueryJobHandle = queryJob.Schedule();
                    mQueryJobCanToComplete = true;
                }
            }
        }

        /// <summary>
        /// 全角色位置信息分帧写入
        /// 信息写入完成后发起KdTree构建任务
        /// </summary>
        public void Update2()
        {
            if (!mIsValid || !Enable || UpdateFrameCount <= 0) return;
            if (!TryGetCamera(out var _)) return;

            if (mQueryJobCanToComplete)
            {
                //完成屏幕内角色查询任务
                mQueryJobCanToComplete = false;
                mQueryJobHandle.Complete();
                mResultBuffer = mResultWrap.Value;
                //写入结果
                mInAreaSet.Clear();
                foreach (var item in mResultBuffer)
                    mInAreaSet.Add((int)item.z);

                //新的一轮开始
                mUpdateFrameLeft = UpdateFrameCount;
                mPositions.Clear();
                if (mTransformTempRemoveCache.Count > 0)
                {
                    foreach (var id in mTransformTempRemoveCache)
                    {
                        mTransformCache.Remove(id);
                        mIdList.Remove(id);
                    }

                    mTransformTempRemoveCache.Clear();
                }

                if (mTransformTempAddCache.Count > 0)
                {
                    foreach (var item in mTransformTempAddCache)
                    {
                        mTransformCache.Add(item.Key, item.Value);
                        mIdList.Add(item.Key);
                    }

                    mTransformTempAddCache.Clear();
                }

                mPerFrameUpdateCount = Mathf.CeilToInt(1f * mIdList.Count / mUpdateFrameLeft);
                if (mPerFrameUpdateCount == 0)
                {
                    mUpdateIdBegin = -1;
                }
                else
                {
                    mUpdateIdBegin = 0;
                    mUpdateIdEnd = mPerFrameUpdateCount - 1;
                }
            }

            mUpdateFrameLeft--;
            if (mUpdateIdBegin >= 0) //合法性检验
            {
                for (int i = mUpdateIdBegin; i <= mUpdateIdEnd; i++) //写入部分数据
                {
                    int id = mIdList[i];
                    var trans = mTransformCache[id];
                    if (trans == null) continue; //防空判断
                    var pos = trans.position;
                    float3 pos3 = new float3(pos.x, pos.z, id);
                    mPositions.Add(pos3);
                }
            }

            if (mUpdateFrameLeft > 0) //更新下一帧待写入数据索引
            {
                int maxId = mIdList.Count - 1;
                if (mUpdateIdEnd >= maxId)
                {
                    mUpdateIdBegin = -1;
                }
                else
                {
                    mUpdateIdBegin = Mathf.Min(mUpdateIdEnd + 1, maxId);
                    mUpdateIdEnd = Mathf.Min(mUpdateIdBegin + mPerFrameUpdateCount - 1, maxId);
                }
            }
            else //数据写入完毕，开始构建KdTree
            {
                mPositionsWrap.Value = mPositions;
                mKdTreeWrap.Value = mKdTree;
                KdTreeBuildJob buildJob = new KdTreeBuildJob()
                {
                    KdTree = mKdTreeWrap,
                    Points = mPositionsWrap,
                };
                mBuildJobHandle = buildJob.Schedule();
                mBuildJobCanToComplete = true;
                IsValidKdTree = false;
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void Add(int id, Transform transform)
        {
            mTransformTempAddCache.Add(id, transform);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void Remove(int id)
        {
            if (!mTransformTempAddCache.Remove(id))
                mTransformTempRemoveCache.Add(id);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool CheckInArea(int id)
        {
            if (!Enable) return true;
            return mInAreaSet.Contains(id);
        }

        public bool TryGetValidKdTree(out UnsafeKdTree2D kdTree)
        {
            if (IsValidKdTree)
            {
                kdTree = mKdTree;
                return true;
            }

            kdTree = default;
            return false;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        bool CalculateArea(out float4 rect)
        {
            if (!TryGetCamera(out var camera))
            {
                rect = default;
                return false;
            }

            //todo 此处需要依据不同项目自行调整
            var ld = camera.ScreenToWorldPoint(new Vector3(0, 0, 0));
            var ru = camera.ScreenToWorldPoint(new Vector3(Screen.width, Screen.height, 0));
            float minX = ld.x;
            float maxX = ru.x;
            float minZ = ld.z + ld.y;
            float maxZ = minZ + 2f * (ru.y - ld.y);
            rect = new float4(
                minX - Tolerance,
                minZ - Tolerance,
                maxX + Tolerance,
                maxZ + Tolerance
            );
            return true;
        }

        bool TryGetCamera(out Camera cam)
        {
            cam = null;
            return false;
        }
    }
}